package com.tree.avl;

import com.tree.BinaryTreeUtils;

import java.util.Scanner;

/**
 * @author `RKC`
 * @date 2021/8/22 10:12
 */
public class AVLTree {

    public static void main(String[] args) {
        TreeNode root = null;
        Scanner scanner = new Scanner(System.in);
        System.out.println("1、添加结点");
        System.out.println("2、删除节点");
        System.out.println("-1、退出");
        int op, val;
        while ((op = scanner.nextInt()) != -1) {
            val = scanner.nextInt();
            switch (op) {
                case 1: root = insert(root, val); break;
                case 2: root = delete(root, val); break;
            }
            BinaryTreeUtils.showTree(root);
            System.out.println("------------------------------");
        }
    }

    /**
     * AVL树的插入操作
     *
     * @param root  根
     * @param value 值
     * @return {@link TreeNode}
     */
    public static TreeNode insert(TreeNode root, int value) {
        if (root == null) return new TreeNode(value);
        if (value > root.value) root.right = insert(root.right, value);
        if (value < root.value) root.left = insert(root.left, value);
        //在递归返回的过程中，对root节点进行平衡
        return balance(root);
    }

    /**
     * AVL树的删除
     *
     * @param root  根
     * @param value 值
     * @return {@link TreeNode}
     */
    public static TreeNode delete(TreeNode root, int value) {
        if (root == null) return null;
        if (root.value < value) {
            //往右子树进行删除
            root.right = delete(root.right, value);
        } else if (root.value > value) {
            //往左子树进行删除
            root.left = delete(root.left, value);
        } else if (root.left != null && root.right != null) {
            //当前root即是删除节点，且当前的root的度为2，找到root的中序前驱节点pre，将前驱节点的value赋值给root，递归向左进行删除pre节点
            TreeNode pre = predecessor(root);
            root.value = pre.value;
            root.left = delete(root.left, pre.value);
        } else {
            //当前被删除节点的度为1或0，temp指针指向唯一子孩子或null，并返回，root没有引用则会被jvm自动回收
            return root.left == null ? root.right : null;
        }
        //递归进行AVL树的平衡操作
        return balance(root);
    }

    /**
     * AVL树的平衡的操作
     *
     * @param cur 当前节点
     * @return {@link TreeNode}
     */
    private static TreeNode balance(TreeNode cur) {
        int leftHeight = height(cur.left);
        int rightHeight = height(cur.right);
        if (leftHeight - rightHeight > 1) {
            //如果当前节点的左子树的左子树高度小于当前节点的左子树的右子树，需要先对当前节点的左子树进行左旋
            if (height(cur.left.left) < height(cur.left.right)) cur.left = leftRotate(cur.left);
            //左子树比右子树高超过1，对当前节点进行右旋
            return rightRotate(cur);
        }
        if (rightHeight - leftHeight > 1) {
            //如果当前节点的右子树的左子树高度大于当前节点的右子树的右子树，需要先对当前结点的右子树进行右旋
            if (height(cur.right.left) > height(cur.right.right)) cur.right = rightRotate(cur.right);
            //右子树比左子树高超过1，对当前节点进行左旋
            return leftRotate(cur);
        }
        //满足平衡条件
        return cur;
    }

    /**
     * 得到root的中序前驱节点
     *
     * @param root 根
     * @return {@link TreeNode}
     */
    private static TreeNode predecessor(TreeNode root) {
        TreeNode pre = root.left;
        while (pre.right != null) pre = pre.right;
        return pre;
    }

    /**
     * 对cur节点进行右旋转
     *
     * @param cur 当前节点
     * @return {@link TreeNode}
     */
    private static TreeNode rightRotate(TreeNode cur) {
        TreeNode temp = cur.left;
        cur.left = temp.right;
        temp.right = cur;
        return temp;
    }

    /**
     * 对cur节点进行左旋转
     *
     * @param cur 当前节点
     * @return {@link TreeNode}
     */
    private static TreeNode leftRotate(TreeNode cur) {
        TreeNode temp = cur.right;
        cur.right = temp.left;
        temp.left = cur;
        return temp;
    }

    /**
     * 求cur节点的高度
     *
     * @param cur 根节点
     * @return int
     */
    public static int height(TreeNode cur) {
        if (cur == null) return 0;
        return 1 + Math.max(height(cur.left), height(cur.right));
    }

    public static class TreeNode {
        public TreeNode left;
        public TreeNode right;
        public int value;

        public TreeNode(int value) {
            this.value = value;
        }

        public TreeNode(int value, TreeNode left, TreeNode right) {
            this(value);
            this.left = left;
            this.right = right;
        }
    }
}
